#include <immintrin.h>
#include <math.h>
#include <stdio.h>
#include "utils.h"

void scalar_nested_conditions(const float *a, const float *b, float *dst, int n)
{
    for (int i = 0; i < n; i++)
    {
        float final = 0;
        float sa = a[i];
        float sb = b[i];
        if (sa > sb)
        {
            final = sa + sb;
        }
        else
        {
            final = sa - sb;
        }
        dst[i] = final;
    }
}

void simd_nested_conditions(const float *a, const float *b, float *dst, int n)
{
    for (int i = 0; i < n; i += 8)
    {
        // 加载数据
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);

        // Step 1: 生成掩码（假设条件为 a > b）
        __m256 mask = _mm256_cmp_ps(va, vb, _CMP_GT_OQ);

        // Step 2: 并行计算所有元素的加法和减法
        __m256 add_result = _mm256_add_ps(va, vb);
        __m256 sub_result = _mm256_sub_ps(va, vb);

        // Step 3: 根据掩码选择结果

        __m256 final = _mm256_blendv_ps(sub_result, add_result, mask);
        // 存储结果
        _mm256_storeu_ps(dst + i, final);
    }
}

int main()
{
    int len = 1;
    for (len = 1000; len < 100000000; len *= 10)
    {
        float *a = aligned_alloc(32, len * sizeof(float));
        float *b = aligned_alloc(32, len * sizeof(float));
        float *dst_scalar = aligned_alloc(32, len * sizeof(float));
        float *dst_simd = aligned_alloc(32, len * sizeof(float));

        // 初始化数据
        for (int i = 0; i < len; ++i)
        {
            a[i] = (float)rand() / RAND_MAX;
            b[i] = (float)rand() / RAND_MAX;
        }
        printf("数据个数 %d\n", len);
        // 标量版本计时
        BENCHMARK_FUNC(10, 100, scalar_nested_conditions, a, b, dst_scalar, len);

        // SIMD版本计时
        BENCHMARK_FUNC(10, 100, simd_nested_conditions, a, b, dst_simd, len);

        // 校验结果一致性
        for (int i = 0; i < len; ++i)
        {
            if (dst_scalar[i] != dst_simd[i])
            {
                printf("dst mismatch at index %d\n", i);
                break;
            }
        }

        free(a);
        free(b);

        free(dst_scalar);
        free(dst_simd);
    }
    return 0;
}